import shutil
import os
import sys
import math
import numpy as np
try:
    from .cvio import cvio
    from . import imagesize
except:
    from cvio import cvio
    import imagesize
from copy import deepcopy
from PIL import Image, ImageDraw
from concurrent.futures import ThreadPoolExecutor, wait, as_completed
import multiprocessing
import time
# import random
# import imagesize


_labelme_format = {
    "imagePath": "",
    "imageHeight": 0,
    "imageWidth": 0,
    "imageData": None,
    "lineColor": [
        0,
        255,
        0,
        128
    ],
    "fillColor": [
        255,
        0,
        0,
        128
    ],
    "flags": {},
    "shapes": [
        # {"line_color":None, "fill_color":None, "label":"",
        #  "shape_type":"polygon", "flags":{}, "points":[], }
    ]
}


def shape_to_mask(
    img_shape, points, shape_type=None, line_width=10, point_size=5
):
    # 根据多边形分割标注轮廓画出并返回目标mask
    mask = np.zeros(img_shape[:2], dtype=np.uint8)
    mask = Image.fromarray(mask)
    draw = ImageDraw.Draw(mask)
    xy = [tuple(point) for point in points]
    if shape_type == "circle":
        assert len(xy) == 2, "Shape of shape_type=circle must have 2 points"
        (cx, cy), (px, py) = xy
        d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
        draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1)
    elif shape_type == "rectangle":
        assert len(xy) == 2, "Shape of shape_type=rectangle must have 2 points"
        draw.rectangle(xy, outline=1, fill=1)
    elif shape_type == "line":
        assert len(xy) == 2, "Shape of shape_type=line must have 2 points"
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == "linestrip":
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == "point":
        assert len(xy) == 1, "Shape of shape_type=point must have 1 points"
        cx, cy = xy[0]
        r = point_size
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
    else:
        assert len(xy) > 2, "Polygon must have points more than 2"
        draw.polygon(xy=xy, outline=1, fill=1)
    mask = np.array(mask, dtype=bool)
    return mask

def get_bbox_from_points(points):
    # 获取多边形的外接矩形框
    points = np.array(points).astype(np.int32)
    xmin = points[:, 0].min()
    xmax = points[:, 0].max()
    ymin = points[:, 1].min()
    ymax = points[:, 1].max()
    return (xmin, ymin, xmax, ymax)


def kernel_crop(img, shapes, with_mask=True, exp_rate=0.2):
    # 根据shape裁剪出图片roi区域
    masks = []
    h, w = img.shape[:-1]
    for shape in shapes:
        label = shape['label']#.split('-')[0]
        points = shape['points']
            
        x1, y1, x2, y2 = get_bbox_from_points(points)
        x1 = max(0, x1)
        x2 = min(x2, w)
        y1 = max(0, y1)
        y2 = min(y2, h)
        box = [x1, y1, x2, y2]
        if exp_rate > 0:
            x1, y1, x2, y2 = expand_bbox([x1, y1, x2, y2], w, h, exp_rate=exp_rate)
        if x1 >= x2 or y1 >= y2 or x2 > w or y2 > h:
            print(label, (h, w), [x1, y1, x2, y2])
            continue

        roi = img[y1:y2, x1:x2]
        if with_mask:
            # 将目标区域叠加带颜色的mask，方便检查
            _mask = shape_to_mask(img.shape, points)
            mask = np.zeros_like(roi, dtype=np.uint8)
            mask[:, :, 2] = _mask[y1:y2, x1:x2] * 255
            roi = (0.7 * roi + 0.3 * mask).astype(np.uint8)

        masks.append((roi, label, box, points))
    return masks

def crop_from_mask(src, dst="", rmtree=False, recursive=False, with_mask=False, exp_rate=0.2, num_workers=32):

    # 裁剪labelme图片标注
    if dst == "":
        dst = os.path.join(src, 'cuts')

    if src == dst:
        print('该目录不允许操作.')
        return 
    
    if os.path.exists(dst):
        print('目录已存在（%s）' % dst)
        # return
        if rmtree:
            shutil.rmtree(dst)

    if not os.path.exists(dst):
        os.makedirs(dst)
        
    imgs_anns = cvio.load_img_ann_list(src, recursive=recursive, ann_type='json')
    n = len(imgs_anns)
    print(src, n)

    annotations = {}
    # m = 0
    # write_ann = cvio.write_ann
    # load_img = cvio.load_img
    # write_img = cvio.write_img
    # load_ann = cvio.load_ann
    # s = time.time()
    # for i, (imgpath, annpath) in enumerate(imgs_anns, 1):
    #     try:
    #         e = time.time()
    #         print('[%d/%d][%.2f%%](%.2fs)' % (i, n, i / n * 100, e - s), os.path.basename(imgpath))
    #         img = np.array(load_img(imgpath, mode='pil'))
    #         h, w = img.shape[:2]
    #         ann_info = load_ann(annpath, ann_type='json')
    #         shapes = ann_info['shapes']
    #         cropped_masks = kernel_crop(img, shapes, with_mask, exp_rate=exp_rate)
    #         # 按类别存放roi图片
    #         if not os.path.exists(dst):
    #             os.makedirs(dst)
    #         imgname = os.path.basename(imgpath)
    #         imgprefix, ext = os.path.splitext(imgname)
    #         for mask in cropped_masks:
    #             mask, label, bbox, points = mask
    #             savepath = os.path.join(dst, label)
    #             if not os.path.exists(savepath):
    #                 os.makedirs(savepath)
    #             m += 1
    #             maskname = '%s%s_%s%s' % (imgprefix, ext, str(bbox), ext)
    #             annotations[maskname] = dict(image=imgname, height=h, width=w, points=points)
    #             maskpath = os.path.join(savepath, maskname)
    #             write_img(mask, maskpath, mode='pil')
    #     except Exception as e:
    #         print(e)
    num_workers = min(num_workers, multiprocessing.cpu_count() * 16)
    
    faster_running = ThreadPoolExecutor(num_workers)
    pools = []
    s = time.time()
    for i, (imgpath, annpath) in enumerate(imgs_anns, 1):
        def cprint(self, i=i, n=n, imgpath=imgpath):
            e = time.time()
            info = '[%d/%d][%.2f%%](%.2fs) %s\n' % (i, n, i / n * 100, e - s, os.path.basename(imgpath))
            sys.stdout.write(info)
        p = faster_running.submit(faster_croping, imgpath, annpath, dst, annotations, with_mask, exp_rate)
        p.add_done_callback(cprint)
        pools.append(p)
    for p in as_completed(pools):
        try:
            p.result()
        except Exception as e:
            print(e)
            faster_running.shutdown(wait=True)

    cvio.write_ann(annotations, os.path.join(dst, 'annotations.json'))
    print('SKU裁剪保存至%s.' % dst)

def faster_croping(imgpath, annpath, dst, annotations={}, with_mask=False, exp_rate=0.2):
    load_img = cvio.load_img
    write_img = cvio.write_img
    load_ann = cvio.load_ann    
    try:
        img = np.array(load_img(imgpath, mode='pil').convert('RGB'))
        h, w = img.shape[:2]
        ann_info = load_ann(annpath, ann_type='json')
        shapes = ann_info['shapes']
        cropped_masks = kernel_crop(img, shapes, with_mask, exp_rate=exp_rate)
        # 按类别存放roi图片

        imgname = os.path.basename(imgpath)
        imgprefix, ext = os.path.splitext(imgname)
        for mask in cropped_masks:
            mask, label, bbox, points = mask
            savepath = os.path.join(dst, label)
            if not os.path.exists(savepath):
                try:
                    os.makedirs(savepath)
                except:
                    pass
            maskname = '%s%s_%s%s' % (imgprefix, ext, str(bbox), ext)
            annotations[maskname] = dict(image=imgname, height=h, width=w, points=points)
            maskpath = os.path.join(savepath, maskname)
            write_img(mask, maskpath, mode='pil')
    except Exception as e:
        print(e)    

def expand_bbox(bbox, W, H, exp_rate=0.2):
    x1, y1, x2, y2 = bbox
    w = x2 - x1
    h = y2 - y1
    w_delta = int(w * exp_rate / 2)
    h_delta = int(h * exp_rate / 2)
    x1 = max(0, x1 - w_delta)
    x2 = min(x2 + w_delta, W)
    y1 = max(0, y1 - h_delta)
    y2 = min(y2 + h_delta, H)

    return (x1, y1, x2, y2)

def gen_mask_from_crop(cropsrc, imgsrc, dst="", recursive=False):
    if dst in (None, "", imgsrc):
        dst = os.path.join(imgsrc, 'new_mask')

    annPath = os.path.join(cropsrc, 'annotations.json')
    if not os.path.exists(annPath):
        print('标注信息损失', annPath, '不存在, 无法恢复标注轮廓信息.')
        annotations = None
        # return
    else:
        annotations = cvio.load_ann(annPath)

    if not os.path.exists(dst):
        os.makedirs(dst)
    
    existsAnns = cvio.load_ext_list(dst)
    if len(existsAnns):
        tmpdir = os.path.join(imgsrc, 'last_mask')
        print('目录（%s）下存在%d个json旧文件，已移存至%s.' % (imgsrc, len(existsAnns), tmpdir))
        if not os.path.exists(tmpdir):
            os.makedirs(tmpdir)
        for t in existsAnns:
            tmpsave = os.path.join(tmpdir, os.path.basename(t))
            if os.path.exists(tmpsave):
                path, ext = tmpsave.split(os.sep)
                tmpsave = '%s_new%s' % (path, ext)
            try:
                shutil.move(t, tmpsave)
            except Exception as e:
                print(e)

    # 加载所有裁剪子图

    cropped_masks = cvio.load_image_list(cropsrc, recursive=True)
    n1 = len(cropped_masks)
    print("收集到%d个SKU" % (n1))
    image_info = {}

    # 收集整理图片的标注信息
    for i,crop in enumerate(cropped_masks,1):
        try:
            if annotations is not None:
                annotation = annotations[os.path.basename(crop)]
            else:
                annotation = get_crop_annotation(imgsrc, crop)
            imgname = annotation['image']
            label = crop.split(os.sep)[-2]
            points = annotation['points']
            shape_type = 'rectangle' if len(points) < 3 else 'polygon' 
            shape = dict(label=label, points=points, shape_type=shape_type)
            if not imgname in image_info.keys():
                ann_info = deepcopy(_labelme_format)
                ann_info['imagePath'] = imgname
                # w, h = imagesize.get(imgpath)
                try:
                    h, w = annotation['height'], annotation['width']
                    assert h > 0 and w > 0
                except:
                    imgpath = os.path.join(imgsrc, imgname)
                    w, h = imagesize.get(imgpath)
                ann_info['imageHeight'] = h
                ann_info['imageWidth'] = w
                ann_info['shapes'].append(shape)
                image_info[imgname] = ann_info
            else:
                image_info[imgname]['shapes'].append(shape)
            if i % 10000 == 0 and i >= 10000:
                print('收集标注[%d/%d] %.2f%%' % (i, n1, i / n1 * 100))
        except Exception as e:
            print(e)
            return

    # 生成图片标注json文件
    imgList = cvio.load_image_list(imgsrc, recursive=recursive)
    n2 = len(imgList)
    for j, imgpath in enumerate(imgList,1):
        imgname = os.path.basename(imgpath)
        print("生成新标注Json [%d/%d]" % (j, n2), imgname)
        if imgname in image_info:
            ann_info = image_info[imgname]
        else:
            ann_info = deepcopy(_labelme_format)
            ann_info['imagePath'] = imgname
        if recursive:
            jsonpath = os.path.splitext(imgname)[0] + '.json'
        else:
            jsonpath = os.path.join(dst, '%s.json' % os.path.splitext(imgname)[0])
        cvio.write_ann(ann_info, jsonpath)

    print('生产新标注结果保存至%s.' % (dst))

def get_crop_annotation(imgsrc, cropsrc):
    imgpath = os.path.join(imgsrc, os.path.basename(cropsrc).split('_[')[0])
    imgname = os.path.basename(imgpath)
    label = cropsrc.split(os.sep)[-2]
    w, h = imagesize.get(imgpath)
    try:
        right = cropsrc.split('_[')[1]
        bbox, ext = right.split('].')
        x1, y1, x2, y2 = eval('[%s]' % bbox)
        points = [[x1, y1], [x2, y2]]
    except Exception as e:
        print(e)
        right = cropsrc.split('_points[')[1]
        points, ext = right.split('].')
        points = eval('[%s]' % (points))
    finally:
        annotation = dict(image=imgname, height=h, width=w, label=label, points=points)
        return annotation



if __name__ == '__main__':
    imgsrc = r'G:/data/datasets/drink/huiyuan/test_llw/test'
    cropsrc = r"G:\data\datasets\drink\huiyuan\test_llw\test_cuts"
    # imgsrc = r'F:\Dataset\jiujiang\modeling_data\labelme\checked\train'
    # cropsrc = r'F:\Dataset\jiujiang\modeling_data\labelme\checked\train_cut'
    # maskdst = r'F:\Dataset\jiujiang\modeling_data\labelme\checked\train'
    maskdst = ''#r'G:/data/datasets/drink/huiyuan/test_llw/test-regen'
    ##  1.裁剪mask，检查标签类别是否错误
    # with_mask：裁剪目标是否带上透明颜色（方便检查）
    
    # crop_from_mask(imgsrc, cropsrc, recursive=True, with_mask=False, exp_rate=0)
    ## 2.检查完以后,重新生成json
    gen_mask_from_crop(cropsrc, imgsrc, maskdst)
